#! /srv01/technion/snirgaz/Software/epd-7.2-2-rh5-x86_64/bin/python2.7

# Imports

import os, sys, glob, math, stat, shutil, subprocess
import xml.etree.ElementTree as etree
import argparse
import random
import csv
import tables as tb
import matplotlib.pylab as plt
import numpy as np
import pickle

class PlotSaveMaxEnt:
    def plotSaveData(self, data, mxData):
        self.mxData = mxData
        analyzedData = self.analyzeData(data)
        self.plotData(analyzedData)        
        return {"mxData":mxData,"anData":analyzedData}        
    def analyzeData(self, data):
        anData = {}
        anData["G"]=self.mxData.G
        anData["MU"]=self.mxData.MU
	anData["beta"]=self.mxData.beta
        anData["alpha"] = [x.alpha for x in data]
        anData["entropy"] = [x.entropy for x in data]
        anData["chisqr"] = [x.chisqr for x in data]
	anData["diff"] = [x.diff for x in data]
        anData["Q"] = [x.Q for x in data]
        # Calc Probability
        probAlpha = [x.probAlpha for x in data]
        probAlpha = np.array(probAlpha)
        probAlpha -= np.max(probAlpha)
        probAlpha = np.exp(probAlpha)
        probAlpha /= np.sum(probAlpha)
        probAlpha = probAlpha.tolist()
        anData["probAlpha"] = probAlpha
        anData["logProbAlpha"] = [math.log(x + 1E-60) for x in probAlpha]
        # Average
        fdata = []
	dw=self.mxData.omegaP[1]-self.mxData.omegaP[0]
        for x in data:
            fo = np.squeeze(np.array(x.f))
            fo = np.multiply(1-np.exp(-self.mxData.omegaP*float(self.mxData.beta)), fo)
	    fo = fo/dw
            fdata.append(fo)
        anData["A_All"] = fdata
        anData["A_Average"] = reduce(lambda x, y:x + y, [p * a for p, a in zip(probAlpha, fdata)])
        anData["A_Smooth"] = fdata[1]
        # max 
        max_pos = np.array(probAlpha).argmax()
        anData["A_maxProb"] = fdata[max_pos]  
        anData["curvature"] = self.curvature(anData["alpha"], anData["entropy"], anData["chisqr"])
        anData["A_maxCurve"] = fdata[anData["curvature"]["maxk"]]
        anData["chisqr_maxA"] = anData["chisqr"][max_pos] 
        anData["chisqr_maxCurve"] = anData["chisqr"][anData["curvature"]["maxk"]] 
        anData["diff_maxA"] = anData["diff"][anData["curvature"]["maxk"]]
        anData["diff_maxCurve"] = anData["diff"][anData["curvature"]["maxk"]]      
	#error
	anData["err"]=[]
	anData["err_pos"]=[]
	for interval in self.mxData.errorIntervals:
		pos=np.where(np.logical_and(self.mxData.omegaP>=interval[0], self.mxData.omegaP<=interval[1]))
		anData["err"].append(np.sum(data[max_pos].nabQ[pos,pos]))
		midpos=pos[0][len(pos[0])/2]
		anData["err_pos"].append(midpos)
        return anData 
    def curvature(self, alpha, entropy, chisqr):
        entropy = np.log(-np.array(entropy))
        chisqr = np.log(np.array(chisqr))
        alpha = np.array(alpha)
        # 1st D
        dalpha = np.diff(alpha)
        dentropy = np.diff(entropy) / dalpha
        dchisqr = np.diff(chisqr) / dalpha
        dalpha = dalpha[0:-1]
        # 2ed D
        d2entropy = np.diff(dentropy) / dalpha
        d2chisqr = np.diff(dchisqr) / dalpha
        dentropy = dentropy[0:-1]
        dchisqr = dchisqr[0:-1]
        # curvature
        k = np.abs((dentropy * d2chisqr - dchisqr * d2entropy) / np.power(dchisqr * dchisqr + d2entropy * d2entropy, 3. / 2.))
        return {"alpha":alpha[0:-2], "k":k, "maxk":k.argmax()}
    def plotData(self, anData):
        fig = plt.figure()
        ax = fig.add_subplot(111)
        ax.plot(self.mxData.omegaP, anData["A_Average"])
        ax.set_xlim(0, 10)
        fig.savefig(self.mxData.dump_dir + "/averageA.pdf")
        fig.clf()
        ax = fig.add_subplot(111)
        ax.plot(self.mxData.omegaP, anData["A_maxProb"])
	omegaerr=[self.mxData.omegaP[p] for p in anData["err_pos"]]
	valerr=[anData["A_maxProb"][p] for p in anData["err_pos"]]
	#print np.sqrt(anData["err"])
	#ax.plot(omegaerr, valerr, yerr=np.sqrt(anData["err"]), fmt='o')
        ax.set_xlim(0, 10)
        fig.savefig(self.mxData.dump_dir + "/maxA.pdf")
        fig.clf()
        ax = fig.add_subplot(111)
        ax.plot(self.mxData.omegaP, anData["A_maxCurve"])
        ax.set_xlim(0, 10)
        fig.savefig(self.mxData.dump_dir + "/maxCurvatureA.pdf")
        fig.clf() 
        ax = fig.add_subplot(111)
        ax.semilogx(anData["alpha"], anData["probAlpha"])
        fig.savefig(self.mxData.dump_dir + "/probAlpha.pdf")
        ax.set_ylim(np.min(anData["probAlpha"]), np.max(anData["probAlpha"]))
        fig.clf()
        ax = fig.add_subplot(111)
        ax.semilogx(anData["alpha"], anData["logProbAlpha"])
        ax.set_ylim(np.min(anData["logProbAlpha"]), np.max(anData["logProbAlpha"]))
        fig.savefig(self.mxData.dump_dir + "/logProbAlpha.pdf")
        fig.clf()
        ax = fig.add_subplot(111)
        ax.semilogx(anData["alpha"], anData["entropy"])
        ax.set_ylim(np.min(anData["entropy"]), np.max(anData["entropy"]))
        fig.savefig(self.mxData.dump_dir + "/entropy.pdf")
        fig.clf()
        ax = fig.add_subplot(111)
        ax.semilogx(anData["alpha"], anData["chisqr"])
        ax.set_ylim(np.min(anData["chisqr"]), np.max(anData["chisqr"]))
        fig.savefig(self.mxData.dump_dir + "/chisqr.pdf")
        fig.clf()
        ax = fig.add_subplot(111)
        ax.semilogx(anData["alpha"], anData["Q"])
        ax.set_ylim(np.min(anData["Q"]), np.max(anData["Q"]))
        fig.savefig(self.mxData.dump_dir + "/Q.pdf")
        fig.clf()
        ax = fig.add_subplot(111)
        ax.plot([math.log(abs(x)) for x in anData["entropy"]], [math.log(abs(x)) for x in anData["chisqr"]])        
        fig.savefig(self.mxData.dump_dir + "/Lcurve.pdf")
        fig.clf()
        ax = fig.add_subplot(111)
        ax.plot(anData["curvature"]["alpha"], anData["curvature"]["k"])        
        fig.savefig(self.mxData.dump_dir + "/curvature.pdf")
        fig.clf()
	ax = fig.add_subplot(111)
        ax.plot(anData["diff_maxCurve"])        
        fig.savefig(self.mxData.dump_dir + "/diffCurve.pdf")
        fig.clf()
	ax = fig.add_subplot(111)
        ax.plot(anData["diff_maxA"])        
        fig.savefig(self.mxData.dump_dir + "/diffMaxEnt.pdf")
        fig.clf()
        ax = fig.add_subplot(111)
        ax.plot(self.mxData.omegaP, anData["A_Smooth"])
        ax.set_xlim(0, 10)
        fig.savefig(self.mxData.dump_dir + "/smoothA.pdf")
        ax = fig.add_subplot(111)
        for a, alpha in zip(anData["A_All"], anData["alpha"]):
            anp = np.transpose(np.squeeze(np.array(a)))
            #ax.plot(self.mxData.omegaP, np.multiply(self.mxData.omegaP, anp))
            ax.plot(self.mxData.omegaP, anp, label='{0:.3f}'.format(alpha))
        ax.set_xlim(0, 10)
        handles, labels = ax.get_legend_handles_labels()
        # reverse the order
        ax.legend(handles[::-1], labels[::-1])
        # or sort them by labels
        import operator
        hl = sorted(zip(handles, labels),
                    key=operator.itemgetter(1))
        handles2, labels2 = zip(*hl) 
        ax.legend(handles2, labels2)
        fig.savefig(self.mxData.dump_dir + "/allA.pdf")
	datfile=open(self.mxData.dump_dir + "/chisqr.dat","w")
	datfile.write("chisqr maxA {0}\n".format(anData["chisqr_maxA"]))
	datfile.write("chisqr maxA {0}\n".format(anData["chisqr_maxCurve"]))
	datfile.close()

